package com.code.ape.codeape.common.elasticsearch;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.AllArgsConstructor;
import lombok.NoArgsConstructor;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * @author 公众号：码猿技术专栏
 * @url: www.java-family.cn
 * @description 封装ElasticSearch的工具方法
 */
@AllArgsConstructor
@NoArgsConstructor
public class ResetElasticSearchClient {

	private RestHighLevelClient client;

	/**
	 * 添加数据
	 *
	 * @param index   索引的名称
	 * @param id      主键ID
	 * @param content 内容
	 * @return {@link IndexResponse}
	 */
	public IndexResponse add(String index, String id, Object content) throws Exception {
		IndexRequest request = new IndexRequest(index)
				.id(id)
				.source(new ObjectMapper().writer().writeValueAsString(content), XContentType.JSON);
		return client.index(request, RequestOptions.DEFAULT);
	}

	/**
	 * 添加数据
	 *
	 * @param index       索引名称
	 * @param id          主键ID
	 * @param jsonContent 内容，json字符串
	 * @return {@link IndexResponse}
	 */
	public IndexResponse add(String index, String id, String jsonContent) throws Exception {
		IndexRequest request = new IndexRequest(index)
				.id(id)
				.source(jsonContent, XContentType.JSON);
		return client.index(request, RequestOptions.DEFAULT);
	}

	/**
	 * 添加数据，ElasticSearch自动生成ID
	 *
	 * @param index   索引的名称
	 * @param content 内容
	 * @return {@link IndexResponse}
	 */
	public IndexResponse add(String index, Object content) throws Exception {
		IndexRequest request = new IndexRequest(index)
				.source(new ObjectMapper().writer().writeValueAsString(content), XContentType.JSON);
		return client.index(request, RequestOptions.DEFAULT);
	}

	/**
	 * 批量添加数据
	 *
	 * @param index        索引的名称
	 * @param ids          id集合
	 * @param jsonContents 内容集合
	 * @return {@link IndexResponse}
	 */
	public BulkResponse addBatch(String index, List<String> ids, List<String> jsonContents) throws Exception {
		BulkRequest request = new BulkRequest();
		for (int i = 0; i < jsonContents.size(); i++) {
			String jsonContent = jsonContents.get(i);
			String id = ids.get(i);
			request.add(new IndexRequest(index).id(id).source(jsonContent, XContentType.JSON));
		}
		return client.bulk(request, RequestOptions.DEFAULT);
	}

	/**
	 * 删除数据
	 *
	 * @param index        索引名称
	 * @param queryBuilder 删除条件
	 * @return {@link BulkByScrollResponse}
	 */
	public BulkByScrollResponse delete(String index, QueryBuilder queryBuilder) throws Exception {
		DeleteByQueryRequest request = new DeleteByQueryRequest(index);
		request.setQuery(queryBuilder);
		return client.deleteByQuery(request, RequestOptions.DEFAULT);
	}

	/**
	 * 根据ID删除数据
	 *
	 * @param index 索引名称
	 * @param id    ID
	 * @return {@link DeleteResponse}
	 */
	public DeleteResponse delete(String index, String id) throws Exception {
		DeleteRequest request = new DeleteRequest(index, id);
		return client.delete(request, RequestOptions.DEFAULT);
	}


	/**
	 * 更新数据
	 *
	 * @param index       索引名称
	 * @param id          ID
	 * @param jsonContent 内容，JSON字符串
	 * @return {@link UpdateResponse}
	 */
	public UpdateResponse update(String index, String id, String jsonContent) throws Exception {
		UpdateRequest request = new UpdateRequest(index, id).doc(jsonContent, XContentType.JSON);
		return client.update(request, RequestOptions.DEFAULT);
	}

	/**
	 * 更新数据
	 *
	 * @param index   索引名称
	 * @param id      ID
	 * @param content 内容
	 * @return {@link UpdateResponse}
	 */
	public UpdateResponse update(String index, String id, Object content) throws Exception {
		UpdateRequest request = new UpdateRequest(index, id).doc(content);
		return client.update(request, RequestOptions.DEFAULT);
	}


	/**
	 * 根据条件查询数据
	 *
	 * @param index         索引
	 * @param startPage     开始页
	 * @param pageSize      每页条数
	 * @param sourceBuilder 查询返回条件
	 * @param queryBuilder  查询条件
	 */
	public SearchResponse searchDataPage(String index, long startPage, long pageSize,
														   SearchSourceBuilder sourceBuilder, QueryBuilder queryBuilder)throws Exception {
		SearchRequest request = new SearchRequest(index);
		//设置超时时间
		sourceBuilder.timeout(new TimeValue(3, TimeUnit.SECONDS));
		//设置是否按匹配度排序
		sourceBuilder.explain(true);
		//加载查询条件
		sourceBuilder.query(queryBuilder);
		sourceBuilder.from((int) ((startPage - 1) * pageSize)).size((int) pageSize);
		request.source(sourceBuilder);
		return client.search(request, RequestOptions.DEFAULT);
	}


	/**
	 * 获取索引下的文档总数
	 * @param index 索引
	 * @param sourceBuilder 查询返回条件
	 * @param queryBuilder 查询条件
	 * @return {@link CountResponse}
	 */
	public CountResponse getTotal(String index,SearchSourceBuilder sourceBuilder, QueryBuilder queryBuilder)throws Exception {
		CountRequest countRequest = new CountRequest(index);
		sourceBuilder.query(queryBuilder);
		countRequest.source(sourceBuilder);
		return client.count(countRequest, RequestOptions.DEFAULT);
	}
}
